from __future__ import division
#Image manipulation functions for pre-processing images before registration.

import Image
from scipy.ndimage import histogram, gaussian_gradient_magnitude, gaussian_filter
import scipy.ndimage as ndi
import numpy as np
from pylab import figure, subplot, imshow, plot
import pylab

def make_pdf(img,normed=True):
    """
    Calculate the normalized probability distribution function for image, img.  
    Return pdf as 1D array.  normed (default = True) will return a normalized 
    distribution instead of pixel counts.
    """
    #make sure image is an array
    img = np.asarray(img,'f')
    #must convert to an 8 bit int image
    img = ConvertTo8BitImage(img)
    
    h = histogram(img,0,254,256)
    if normed:
        pdf = h/h.sum()
        return pdf
    else:
        return h

def make_cdf(img):
    """
    Calculate the cumulative distribution function for image, img. 
    Return cdf as 1D array
    """
    #Get the pdf first 
    pdf = make_pdf(img)
    cdf = np.cumsum(pdf)
    return cdf

def histeq(img):
    """
    Perform a histogram equalization on image, img.  Return equalized image.
    """
    img = np.asarray(img,'f')
    #must convert to an 8 bit int image
    img = ConvertTo8BitImage(img)
    
    cdf = make_cdf(img)
    img = 255*cdf[img]
    img = img.astype('int')
    return img

def plot_histeq(img):
    """
    Plot a before and after image grid
    """
    pdf1=make_pdf(img)
    cdf1=make_cdf(img)
    
    J = histeq(img)
    pdf2 = make_pdf(J)
    cdf2 = make_cdf(J)
    
    pylab.figure()
    pylab.subplot(221)
    pylab.imshow(img)
    pylab.subplot(222)
    pylab.plot(pdf1*1/pdf1.max())
    pylab.plot(cdf1)
    pylab.subplot(223)
    pylab.imshow(J)
    pylab.subplot(224)
    pylab.plot(pdf2*1/pdf2.max())
    pylab.plot(cdf2)
 
def edge(img, sigma=1):
    """
    Return an edge image
    Created using the gaussian_gradient_magnitude filter in ndimage.
    """ 
    img = make_gray(img)
    img = gaussian_gradient_magnitude(img,sigma)    
    return img

def enhance_edges(img,sigma=0.5,threshold=0.5):
    """
    Return and edge enhanced image.  This filter will accentuate the
    visible edges in an image by multiplying the original image by the 
    edge image.  The two parameters are:
        sigma - amount of blurring (standard deviations of gaussian blur)
        threshold - to discard spurious edges (0: all edges used, 1: no edges used)
    """
    # only works with 8bit images currently (only applies for threshold)
    #img = ConvertTo8BitImage(img)
    
    # First we create an edge image
    img_edge = edge(img,sigma)
    
    # Take the inverse of the image
    img_edge = img_edge.max() - img_edge
    
    # Discard edge below threshold
    # A cummulative density function is used to throw away a % of unwanted edges
    #cdf_edge = make_cdf(img_edge)
    #thresh_index = cdf_edge.searchsorted(threshold)
    #print "Threshold index: %i" % thresh_index
    #img_edge[img_edge<thresh_index]=0
    
    # enhange the edges in original image
    img_edge = img_edge * img
    
    # Equalize the histogram
    img_edge = histeq(img_edge)
    
    # return enhanced image!
    return img_edge

def blur(img, sigma=1):
    """
    Return a blurred image.
    Created using a gaussian filter with standard deviation, sigma
    set the same in all dimensions.
    """
    img = make_gray(img)
    img = gaussian_filter(img,sigma)
    return img

def make_gray(img):
    """
    Input either a PIL or numpy array.
    
    Return a grayscale intensity image (numpy array).
    """
    if hasattr(img,'convert'):
        print "Converting from PIL to numpy array"
        img = img.convert('L')
        img = np.asarray(img,'f')
    elif hasattr(img,'ndim'):
        img = np.asarray(img,'f')
        if img.ndim > 2:
            # If the image is already a numpy array, try to convert it
            # if ndims is > 2
            print "converting numpy array to grayscale"
            img = img[:,:,0]*.299 + img[:,:,1]*.587 + img[:,:,2]*.114
    else:
        print "image type not recognized as PIL or numpy, exiting"
        return
    
    return img

def ConvertTo8BitImage(img):
    """
    Scale an image to 0-255 integers of uint8
    """
    img = np.asarray(img,'f')
    scaledImage = (img-img.min())/(img.max()-img.min())*255
    return scaledImage.astype('uint8')

def ThresholdImage(img,threshold=None):
    """
    Scale an img to 0-255 integers of uint8
    """
    img = np.asarray(img,'f')
    
    if threshold == None:
        threshold = 0.5*(img.max()-img.min())
    
    img[img<threshold] = 0
    img[img>threshold] = 255
    return img.astype('uint8')

def EdgeDiagnostics(img1,img2,et=0):
    """
    Try to find optimum edge detection on two images and display results
    """
    # Convert images to numpy array of floats
    img1 = np.asarray(img1, 'f')
    img2 = np.asarray(img2, 'f')
    
    # Improve contrast with histeq function
    #img1 = histeq(img1)
    #img2 = histeq(img2)
    
    # Use sobel operator to find edges
    edge1 = ndi.generic_gradient_magnitude(img1.copy(), ndi.sobel)
    edge2 = ndi.generic_gradient_magnitude(img2.copy(), ndi.sobel)
    
    # Threshold out weaker edges using et number of standard deviations
    edge1_thresh = edge1.copy()
    edge2_thresh = edge2.copy()
    
    edge1_thresh[edge1<et*edge1.std()] = 0
    edge2_thresh[edge2<et*edge2.std()] = 0
    
    # Calculate the morphological gradient
    #edge1_mg = ndi.morphological_gradient(img1, 5)
    
    # Display results in a new figure
    figure()
    subplot(321)
    imshow(img1)
    subplot(322)
    imshow(img2)
    subplot(323)
    plot(make_pdf(edge1),hold=True)
    #plot(make_pdf(edge1_thresh))
    subplot(324)
    plot(make_pdf(edge2),hold=True)
    #plot(make_pdf(edge2_thresh))
    subplot(325)
    imshow(edge1_thresh)
    subplot(326)
    imshow(edge2_thresh)
    figure()
    imshow(edge1_thresh-edge2_thresh)


    
    